package com.raos.example.memory;

import com.raos.example.memory.entity.MyChatMessage;
import com.raos.example.memory.repository.MyChatMessageRepository;
import dev.langchain4j.community.model.zhipu.ZhipuAiStreamingChatModel;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * 会话聊天持久化Controller
 *
 * @author raos
 * @date 2025-07-12 16:50
 */
@Slf4j
@RequestMapping("/memory")
@RestController
public class MemoryController {

    @Resource(name = "zhipuAiStreamingChatModel")
    private ZhipuAiStreamingChatModel zhipuAiStreamingChatModel;
    @Autowired
    private MyChatMessageRepository chatMessageRepository;

    /**
     * 智普AI-简单流式响应（浏览器输出）
     */
    @RequestMapping("/zhipu")
    public Flux<String> zhipu(
            @RequestParam(value = "memoryId") String memoryId,
            @RequestParam(value = "message") String message,
            HttpServletResponse response) {
        log.info("memoryId=[{}], message=[{}]", memoryId, message); // 写一篇关于春天的200字的小作文

        // 添加系统角色
        SystemMessage systemMessage = SystemMessage.systemMessage("""
                请以诗人毛泽东的思维来回答用户的问题，如果发现用户提的问题不是诗歌范畴，直接回复：
                您的问题不是诗歌范畴请重新提问
                """);

        List<MyChatMessage> messageList = chatMessageRepository.findAllByChatMemoryId(memoryId);

        List<ChatMessage> chatMessages = new ArrayList<>();
        chatMessages.add(systemMessage);

        // 添加数据库中的问答内容（本质就是上下文）
        for (MyChatMessage message1 : messageList) {
            ChatMessage chatMessage = message1.toChatMessage();
            if (Objects.nonNull(chatMessage)) {
                chatMessages.add(chatMessage);
            }
        }

        MyChatMessage myChatMessage = MyChatMessage.create(memoryId, "USER", message);
        myChatMessage = chatMessageRepository.save(myChatMessage);

        chatMessages.add(UserMessage.userMessage(myChatMessage.getContent()));

        response.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE + "; charset=utf-8");
        return Flux.create(fluxSink -> zhipuAiStreamingChatModel.chat(chatMessages, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String partialResponse) {
                // 每次响应的数据
                //log.info("***AI对话 onNext:\n {}", partialResponse);
                fluxSink.next(partialResponse);
            }

            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                // 响应结束的回调
                log.info("***AI对话 onComplete:\n {}", completeResponse.aiMessage().text());
                MyChatMessage chatMessage = MyChatMessage.create(memoryId, "AI", completeResponse.aiMessage().text());
                chatMessageRepository.save(chatMessage);
                fluxSink.complete();
            }

            @Override
            public void onError(Throwable error) {
                // 响应出错的回调
                log.error("***AI对话，发送异常", error);
                fluxSink.error(error);
            }
        }));
    }

}
